#include "stdafx.h"
#include "CMASearcher.h"
#include "cmaes_interface.h"
#include <stdio.h>
#include "SearcherConfig.h"
#include "ConfigManager.h"
#include "SceneBuilder.h"
#include "Scene.h"
#include "DecoArticulatedBody.h"
#include "timer.h"
#include "DecoLogger.h"
#include "CMASearchContinuation.h"
#include <tbb/tbb.h>
#include <tbb/mutex.h>
#include <tbb/task_scheduler_init.h>
using namespace tbb;


#define NUM_CMA_SEARCH_THREADS 4

extern double (*gEvaluator)(double*, DecoScene*, double*);

class ParallelCMAEvaluator
{
public:
	static tbb::mutex msMutex;

	ParallelCMAEvaluator(CMASearcher* searcher, DecoScene** scenes, double* const* pop, double* arFuncvals, int popSize, int dim) : mSearcher(searcher), mScenes(scenes), mArFuncvals(arFuncvals), mPopSize(popSize), mDim(dim)
	{
		tbb::mutex::scoped_lock lock(msMutex);
		mPop = Malloc2DArray<double>(mPopSize, mDim);
		for (int i = 0; i < mPopSize; ++i)
		{
			for (int j = 0; j < mDim; ++j)
			{
				mPop[i][j] = pop[i][j];
			}
		}
	}
	ParallelCMAEvaluator(const ParallelCMAEvaluator& rhs) : mSearcher(rhs.mSearcher), mScenes(rhs.mScenes), mArFuncvals(rhs.mArFuncvals), mDim(rhs.mDim), mPopSize(rhs.mPopSize)
	{
		tbb::mutex::scoped_lock lock(msMutex);
		mPop = Malloc2DArray<double>(mPopSize, mDim);
		for (int i = 0; i < mPopSize; ++i)
		{
			for (int j = 0; j < mDim; ++j)
			{
				mPop[i][j] = rhs.mPop[i][j];
			}
		}
	}
	void operator() (const blocked_range<size_t>& r) const
	{
		msMutex.lock();
		int ithThread = msIthThread % NUM_CMA_SEARCH_THREADS;
		msIthThread++;
		msMutex.unlock();

		for (size_t i = r.begin(); i != r.end(); ++i)
		{
			if(i != 12)
				continue;
#define RESULT_DETAIL_NUM 10
			double resultDetails[RESULT_DETAIL_NUM];
			DecoTimer timer;
			timer.startTimer();
			mArFuncvals[i] = (*gEvaluator)(mPop[i], mScenes[ithThread], resultDetails);
			resultDetails[RESULT_DETAIL_NUM - 1] = timer.getCurrentTime();


			//try
			//{
			//	mArFuncvals[i] = (*gEvaluator)(mPop[i], mScenes[ithThread]);
			//}
			//catch (CException* e)
			//{
			//	char errorMsg[255];
			//	e->GetErrorMessage(errorMsg, 255);
			//	cmaes_WriteLogFile(errorMsg);
			//	cmaes_WriteLogFile("\n");
			//}
			tbb::mutex::scoped_lock lock(msMutex);

			cmaes_WriteLogFile("(");
			char logStr[512];
			memset(logStr, 0, 512 * sizeof(char));
			sprintf(logStr, "Current Q: (%f, %f, %f, %f, %f, %f) and energy consumption %f and collision %f, took %f seconds\n", resultDetails[1], resultDetails[2], resultDetails[3], resultDetails[4], resultDetails[5], resultDetails[6], resultDetails[7], resultDetails[8], resultDetails[9]);
			cmaes_WriteLogFile(logStr);
			cmaes_WriteLogFile("(");
			for (int j = 0; j < mDim; ++j)
			{
				memset(logStr, 0, 512 * sizeof(char));
				sprintf(logStr, " %f, ", mPop[i][j]);
				cmaes_WriteLogFile(logStr);

			}
			memset(logStr, 0, 512 * sizeof(char));
			sprintf(logStr, "):   %f\n", mArFuncvals[i]);
			cmaes_WriteLogFile(logStr);

		}
	}
	~ParallelCMAEvaluator()
	{
		tbb::mutex::scoped_lock lock(msMutex);
		Free2DArray<double>(mPop, mPopSize, mDim);
	}

private:
	double* mArFuncvals;
	double** mPop;
	DecoScene** mScenes;
	int mDim;
	int mPopSize;
	CMASearcher* mSearcher;
	static int msIthThread;
};

tbb::mutex ParallelCMAEvaluator::msMutex;
int ParallelCMAEvaluator::msIthThread = 0;

CMASearcher::CMASearcher() : mDim(-1), mAssigned(-1), mPrevSol(NULL), mStandardDeviation(NULL), mInfeasibleTime(NULL)
{

}
CMASearcher::CMASearcher(int dim, int assigned) : mDim(-1), mAssigned(-1), mPrevSol(NULL), mStandardDeviation(NULL), mInfeasibleTime(NULL)
{
	mAssigned = assigned;
	SetDimension(dim - assigned);
}
CMASearcher::~CMASearcher()
{
	clear();

}

void CMASearcher::SetDimension(int dim)
{
	clear();
	mDim = dim;
	mPrevSol = new double[mDim];
	mStandardDeviation = new double[mDim];
	mInfeasibleTime = new int[mDim];
}

void CMASearcher::clear()
{
	if (mPrevSol)
		delete[] mPrevSol;
	if (mStandardDeviation)
		delete[] mStandardDeviation;
	if (mInfeasibleTime)
		delete[] mInfeasibleTime;
}

void CMASearcher::setInitialGuess(double* lower_bound, double* upper_bound)
{
	for (int i = 0; i < mDim; ++i)
	{
		mPrevSol[i] = 0.5 * (lower_bound[i] + upper_bound[i]);
	}
	calculateSearchStandardDeviation(lower_bound, upper_bound);

}

// void CMASearcher::SetSineParams(const SearcherInfo& info, double* pop)
// {
// 	for(int i = 0; info.mJointBounds.size(); i++)
// 	{
// 		const JointBound& bound = info.mJointBounds[i];
// 		if(!strncmp(bound.mRawName.c_str(), "min", 3))
// 		{
// 			double dMin = pop[i];
// 			double dMax = pop[i + 1];
// 
// 			pop[i]
// 		}
// 	}
// }

int CMASearcher::RestartSearch(double* lower_bound, double* upper_bound, double* argMin, int maxIterations, int* pSelectedNodes)
{
	setInitialGuess(lower_bound, upper_bound);
	return Search(lower_bound, upper_bound, argMin, maxIterations, pSelectedNodes);
}



int CMASearcher::Search(double* lower_bound, double* upper_bound, double* argMin, int maxIterations, int* pSelectedNodes)
{
	setInitialGuess(lower_bound, upper_bound);

	int iterationCount = 0;

	cmaes_t evo; /* an CMA-ES type struct or "object" */
	double *arFunvals, *const*pop, *xfinal;
	int i; 
	const SearcherInfo& info = DecoConfig::GetSingleton()->GetSearcherInfo();
	/* Initialize everything into the struct evo, 0 means default */
	arFunvals = cmaes_init(&evo, mDim, mPrevSol, mStandardDeviation, 0, 0, info.mCMAInitialPath.c_str()); 
	cmaes_WriteLogFile(cmaes_SayHello(&evo));
	cmaes_WriteLogFile("\n");
	cmaes_ReadSignals(&evo, info.mCMASignalPath.c_str());  /* write header and initial values */
	DecoScene* scenes[NUM_CMA_SEARCH_THREADS];
	for (int ithScene = 0; ithScene < NUM_CMA_SEARCH_THREADS; ++ithScene)
	{
		scenes[ithScene] = new DecoScene();
		SceneBuilder::GetSingleton()->ConstructScene(scenes[ithScene]);	

	}
	/* Iterate until stop criterion holds */
	tbb::simple_partitioner partitioner;

	while(!cmaes_TestForTermination(&evo))
	{ 

		/* generate lambda new search points, sample population */
		pop = cmaes_SamplePopulation(&evo); /* do not change content of pop */

		 for (i = 0; i < cmaes_Get(&evo, "popsize"); ++i) 
		 {
			 while (!isFeasible(pop[i], lower_bound, upper_bound, info.mJointConstraints))
				 cmaes_ReSampleSingle(&evo, i);
		 }
		 int numSamples = cmaes_Get(&evo, "lambda");

		 DecoTimer timer;
		 timer.startTimer();
		 int grainSize = ceil(static_cast<double>(numSamples) / NUM_CMA_SEARCH_THREADS);

		 parallel_for(blocked_range<size_t>(0, numSamples, grainSize), ParallelCMAEvaluator(this, scenes, pop, arFunvals, numSamples, mDim), partitioner);	
		 double currentTime = timer.getCurrentTime();
		 // for (i = 0; i < cmaes_Get(&evo, "lambda"); ++i) {
		 // 	char logStr[512];
		 // 	arFunvals[i] = (*gEvaluator)(pop[i], scenes[i]);
		 // 	//			arFunvals[i] = Evaluator((pop[i]));
		 // 	cmaes_WriteLogFile("(");
		 // 	for (int j = 0; j < mDim; ++j)
		 // 	{
		 // 		 memset(logStr, 0, 512 * sizeof(char));
		 // 		 sprintf(logStr, " %f ", pop[i][j]);
		 // 		 cmaes_WriteLogFile(logStr);

		 // 	}
		 // 	memset(logStr, 0, 512 * sizeof(char));
		 // 	sprintf(logStr, "):   %f\n", arFunvals[i]);
		 // 	cmaes_WriteLogFile(logStr);
		 // }

		/* update the search distribution used for cmaes_SampleDistribution() */
		cmaes_UpdateDistribution(&evo, arFunvals); 

		double* covar = cmaes_GetNew(&evo, "diag(C)");
		double* mean = cmaes_GetNew(&evo, "xmean");
		double* best = cmaes_GetNew(&evo, "xbestever");
		if (covar && mean && best)
		{
			char covarChar[512], meanChar[512], bestChar[512];
			memset(covarChar, 0, 512 * sizeof(char));
			memset(meanChar, 0, 512 * sizeof(char));
			memset(bestChar, 0, 512 * sizeof(char));
			for (int ithDim = 0; ithDim < mDim; ++ithDim)
			{
				sprintf(covarChar, "%s%f  ", covarChar, covar[ithDim]);
				sprintf(meanChar, "%s%f  ", meanChar, mean[ithDim]);
				sprintf(bestChar, "%s%f  ", bestChar, best[ithDim]);
			}
			cmaes_WriteLogFile("Diag Covariance:   ");
			cmaes_WriteLogFile(covarChar);
			cmaes_WriteLogFile("\n");
			cmaes_WriteLogFile("Mean:             ");
			cmaes_WriteLogFile(meanChar);
			cmaes_WriteLogFile("\n");
			cmaes_WriteLogFile("Best x so far:    ");
			cmaes_WriteLogFile(bestChar);
			cmaes_WriteLogFile("\n");

			free(covar);
			free(mean);
			free(best);
		}
		char logStr[512];

		iterationCount++;
		memset(logStr, 0, 512 * sizeof(char));
		sprintf(logStr, "-------------------%d CMA iteration finished in %f seconds---------------\n", iterationCount, currentTime);
		cmaes_WriteLogFile(logStr);
		cmaes_WriteLogFile("\n\n");

		if (iterationCount >= maxIterations)
		{
			break;
		}
	//	/* read instructions for printing output or changing termination conditions */ 
		cmaes_ReadSignals(&evo, info.mCMASignalPath.c_str());   
		fflush(stdout); /* useful in MinGW */
	}
	if (iterationCount < maxIterations)
	{
		const char* terminationText = cmaes_TestForTermination(&evo);
		cmaes_WriteLogFile("Stop:\n");
		cmaes_WriteLogFile(terminationText); /* print termination reason */
	}

	/* get best estimator for the optimum, xmean */
	xfinal = cmaes_GetNew(&evo, "xbestever"); /* "xbestever" might be used as well */
	cmaes_exit(&evo); /* release memory */ 
	memcpy(mPrevSol, xfinal, mDim * sizeof(double));
	memcpy(argMin, xfinal, mDim * sizeof(double));
	/* do something with final solution and finally release memory */
	free(xfinal); 
	for (int ithScene = 0; ithScene < NUM_CMA_SEARCH_THREADS; ++ithScene)
	{
		delete(scenes[ithScene]);
	}
	//if (strstr(terminationText, "ConditionNumber"))
	//{
	//	return 0;
	//}
	return 1;
}

void CMASearcher::calculateSearchStandardDeviation(double* lower_bound, double* upper_bound)
{
	for (int i = 0; i < mDim; ++i)
	{
		double minBoundary = abs(mPrevSol[i] - lower_bound[i]);
		double maxBoundary = abs(mPrevSol[i] - upper_bound[i]);
		mStandardDeviation[i] = max(minBoundary, maxBoundary) / 2.0;
	}
}

bool CMASearcher::isFeasible(double* value, double* lower_bound, double* upper_bound)
{
	for (int i = 0; i < mDim; ++i)
		if (value[i] < lower_bound[i] || value[i] > upper_bound[i])
		{

			//printf("Infeasible sample hit!\n");
			return false;
		}

	return true;
}


bool CMASearcher::isFeasible(double* value, double* lower_bound, double* upper_bound, const vector<JointSearchConstraint>& constriants)
{
	for (int i = 0; i < mDim; ++i)
	{
		if (value[i] < lower_bound[i] || value[i] > upper_bound[i])
		{

			//printf("Infeasible sample hit!\n");
			return false;
		}
	}

	for(int i = 0; i < constriants.size(); i++)
	{
		const JointSearchConstraint& constraint = constriants[i];
		if(constraint.mType == JSCT_Less)
		{
			if(value[constraint.mConstraintedJointIdx - mAssigned] >= value[constraint.mConstraintToJointIdx - mAssigned])
				return false;
		}
		else if(constraint.mType == JSCT_Greater)
		{
			if(value[constraint.mConstraintedJointIdx - mAssigned] <= value[constraint.mConstraintToJointIdx - mAssigned])
				return false;
		}
	}
	
	return true;
}

//void CalculateAmplitudeRampup(double* inputAmplitude, double* outputAmplitude, int num, double currentTime, double totalRampupTime)
//{
//	double ratio = 1.0;
//	if (currentTime <= totalRampupTime)
//		ratio = currentTime / totalRampupTime;
//	for (int i = 0; i < num; ++i)
//	{
//		outputAmplitude[i] = inputAmplitude[i] * ratio;
//	}
//}
//
//







